package com.huawei.admins.platform.gateway.core.component.filter.factory;

import com.huawei.admins.platform.gateway.core.component.filter.FilterDefinition;
import com.huawei.admins.platform.gateway.core.component.filter.GatewayFilter;
import com.huawei.admins.platform.gateway.core.component.limiter.KeyResolver;
import com.huawei.admins.platform.gateway.core.component.limiter.RateLimiter;
import com.huawei.admins.platform.gateway.core.component.limiter.RedisRateLimiter;
import com.huawei.admins.platform.gateway.core.component.limiter.ServiceRateLimiter;
import com.huawei.admins.platform.gateway.core.component.router.Router;
import com.huawei.admins.platform.gateway.core.config.GatewayProperties;
import com.huawei.admins.platform.gateway.core.util.ServerWebExchangeUtils;
import com.huawei.admins.platform.gateway.support.HttpStatusHolder;
import jakarta.validation.constraints.Min;
import org.springframework.core.style.ToStringCreator;
import org.springframework.http.HttpStatus;
import org.springframework.util.Assert;

import java.util.Map;

import static com.huawei.admins.platform.gateway.core.util.ServerWebExchangeUtils.setResponseStatus;

public class RequestRateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<RequestRateLimiterGatewayFilterFactory.CompositeConfig> {

    public static final String EMPTY_KEY = "____EMPTY_KEY__";

    private final RateLimiter defaultRateLimiter;

    private final KeyResolver defaultKeyResolver;

    private boolean denyEmptyKey = true;

    private String emptyKeyStatusCode = HttpStatus.FORBIDDEN.name();

    public RequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter, KeyResolver defaultKeyResolver, GatewayProperties gatewayProperties) {
        super(gatewayProperties);
        this.defaultRateLimiter = defaultRateLimiter;
        this.defaultKeyResolver = defaultKeyResolver;
    }

    @Override
    public GatewayFilter apply(String routeId, FilterDefinition definition) {

        Map<String, String> filterArgs = definition.getArgs();

        KeyResolver resolver = this.defaultKeyResolver;
        RateLimiter limiter = this.defaultRateLimiter;

        boolean denyEmpty = getOrDefault(filterArgs, "deny-empty-key", this.denyEmptyKey);

        HttpStatusHolder denyStatus = HttpStatusHolder.parse(getOrDefault(filterArgs, "deny-status-code", HttpStatus.TOO_MANY_REQUESTS.name()));
        HttpStatusHolder emptyKeyStatus = HttpStatusHolder.parse(getOrDefault(filterArgs, "empty-key-status-code", this.emptyKeyStatusCode));

        CompositeConfig config = loadConfiguration(routeId);

        return (exchange, chain) -> resolver.resolve(exchange).defaultIfEmpty(EMPTY_KEY).flatMap(key -> {
            if (EMPTY_KEY.equals(key)) {
                if (denyEmpty) {
                    setResponseStatus(exchange, emptyKeyStatus);
                    return exchange.getResponse().setComplete();
                }
                return chain.filter(exchange);
            }
            String currentRouteId = routeId;
            if (currentRouteId == null) {
                Router route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
                if (route != null) {
                    currentRouteId = route.getId();
                }
            }
            return limiter.isAllowed(currentRouteId, key, config).flatMap(response -> {

                for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
                    exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
                }

                if (response.isAllowed()) {
                    return chain.filter(exchange);
                }

                setResponseStatus(exchange, denyStatus);
                return exchange.getResponse().setComplete();
            });
        });
    }

    public KeyResolver getDefaultKeyResolver() {
        return defaultKeyResolver;
    }

    public RateLimiter getDefaultRateLimiter() {
        return defaultRateLimiter;
    }

    public boolean isDenyEmptyKey() {
        return denyEmptyKey;
    }

    public void setDenyEmptyKey(boolean denyEmptyKey) {
        this.denyEmptyKey = denyEmptyKey;
    }

    public String getEmptyKeyStatusCode() {
        return emptyKeyStatusCode;
    }

    public void setEmptyKeyStatusCode(String emptyKeyStatusCode) {
        this.emptyKeyStatusCode = emptyKeyStatusCode;
    }

    @Override
    protected CompositeConfig buildConfig(String routeId, Map<String, String> filterArgs) {

        String localLimiterKey = ServiceRateLimiter.CONFIGURATION_PROPERTY_NAME;
        String redisLimiterKey = RedisRateLimiter.CONFIGURATION_PROPERTY_NAME;

        CompositeConfig compositeConfig = new CompositeConfig();

        compositeConfig.setLocalConfig(new Config()
                .setBurstCapacity(Integer.parseInt(filterArgs.getOrDefault(String.format("%s.burstCapacity", localLimiterKey), "1")))
                .setReplenishRate(Integer.parseInt(filterArgs.getOrDefault(String.format("%s.replenishRate", localLimiterKey), "1")))
                .setRequestedTokens(Integer.parseInt(filterArgs.getOrDefault(String.format("%s.requestedTokens", localLimiterKey), "1"))));

        compositeConfig.setRedisConfig(new Config()
                .setBurstCapacity(Integer.parseInt(filterArgs.getOrDefault(String.format("%s.burstCapacity", redisLimiterKey), "1")))
                .setReplenishRate(Integer.parseInt(filterArgs.getOrDefault(String.format("%s.replenishRate", redisLimiterKey), "1")))
                .setRequestedTokens(Integer.parseInt(filterArgs.getOrDefault(String.format("%s.requestedTokens", redisLimiterKey), "1"))));

        return compositeConfig;
    }

    public static class CompositeConfig {

        private Config localConfig;

        private Config redisConfig;

        public Config getLocalConfig() {
            return localConfig;
        }

        public void setLocalConfig(Config localConfig) {
            this.localConfig = localConfig;
        }

        public Config getRedisConfig() {
            return redisConfig;
        }

        public void setRedisConfig(Config redisConfig) {
            this.redisConfig = redisConfig;
        }
    }

    public static class Config {

        @Min(1)
        private int replenishRate;

        @Min(0)
        private int burstCapacity = 1;

        @Min(1)
        private int requestedTokens = 1;

        public int getReplenishRate() {
            return replenishRate;
        }

        public Config setReplenishRate(int replenishRate) {
            this.replenishRate = replenishRate;
            return this;
        }

        public int getBurstCapacity() {
            return burstCapacity;
        }

        public Config setBurstCapacity(int burstCapacity) {
            Assert.isTrue(burstCapacity >= this.replenishRate, "BurstCapacity(" + burstCapacity
                    + ") must be greater than or equal than replenishRate(" + this.replenishRate + ")");
            this.burstCapacity = burstCapacity;
            return this;
        }

        public int getRequestedTokens() {
            return requestedTokens;
        }

        public Config setRequestedTokens(int requestedTokens) {
            this.requestedTokens = requestedTokens;
            return this;
        }

        @Override
        public String toString() {
            return new ToStringCreator(this).append("replenishRate", replenishRate)
                    .append("burstCapacity", burstCapacity).append("requestedTokens", requestedTokens).toString();

        }

    }
}
